import numpy as np
import heapq
import matplotlib.pyplot as plt
def heuristic(a, b):
    return abs(a[0] - b[0]) + abs(a[1] - b[1])
def a_star_search(maze, start, end):
    all_path = []
    frontier = []
    heapq.heappush(frontier, (0, start))
    came_from = {start: None}
    cost_so_far = {start: 0}
    while frontier:
        _, current = heapq.heappop(frontier)
        if current == end:
            break
        for next in [(current[0] + 1, current[1]), (current[0] - 1, current[1]),
                     (current[0], current[1] + 1), (current[0], current[1] - 1)]:
            if 0 <= next[0] < len(maze) and 0 <= next[1] < len(maze[0]) and maze[next[0]][next[1]] == 0:
                new_cost = cost_so_far[current] + 1
                if next not in cost_so_far or new_cost < cost_so_far[next]:
                    cost_so_far[next] = new_cost
                    priority = new_cost + heuristic(end, next)
                    heapq.heappush(frontier, (priority, next))
                    came_from[next] = current
                    all_path.append([(current[0], current[1]), (next[0], next[1])])
    return came_from, cost_so_far, all_path


# 可视化迷宫
def visualize_maze(maze, path=None, all_path=None):
    fig, ax = plt.subplots()
    def update(fp):
        ax.clear()
        ax.imshow(maze, cmap=plt.cm.binary, interpolation='nearest')
        if path:
            path_x, path_y = zip(*path)
            ax.plot(path_y, path_x, 'gray', linewidth=2)
        if all_path:
            for idx, segment in enumerate(fp):
                x, y = zip(*segment)
                line, = ax.plot(y, x, color=f'C{idx}', linewidth=2)
        plt.draw()
        plt.pause(0.1)

    for i in range(1, len(all_path)):
        update(all_path[:i])
    # update(all_path)


if __name__ == '__main__':
    maze = np.loadtxt('maze1.txt', dtype=int)
    start = (0, 0)
    end = (len(maze) - 1, len(maze[0]) - 1)

    came_from, _, all_path = a_star_search(maze, start, end)
    path = []
    current = end
    while current != start:
        path.append(current)
        current = came_from[current]
    path.append(start)
    path.reverse()
    visualize_maze(maze, path, all_path)
